package com.webcamtracker.color.segmentation;

import com.webcamtracker.image.Color;
import com.webcamtracker.tracker.TrackHistory;
import com.webcamtracker.utils.MathFunctions;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.SortedSet;

public class SegmentationImageTracker implements Tracker {

    //Processors
    private final SearchWindowProvider searchWindowProvider;
    private final List<? extends PreProcessor> preprocessors;
    private final ColorSegmentator segmentator;
    //Target Data
    private BufferedImage targetImage;
    private SortedSet<Segment> targetSegments;
    private Map<Integer, Double> targetDistances;
    private TrackingObjectImpl initialElement;

    private TrackHistory<TrackingObject> history;
    private String targetName;
    private static final String ZOOM = "zoom";
    private static final String SEGMENTS_BY_ID = "segments";

    private boolean zoom_enable = true;
    private static final int HISTORY_SIZE = 5;


    public SegmentationImageTracker() {
        preprocessors = Arrays.asList(new BlurPreProcessor(), new HSVColorPlannerPreProcessor());
        searchWindowProvider = new SearchWindowProviderImpl();
        segmentator = new ColorSegmentator();
    }

    public void arrivedFrame(BufferedImage image) {
        if (targetImage != null) {

            final long now = System.currentTimeMillis();
            //Calculate the search window from the history.
            //Pop the search window and cut a new sub-image.
            final SearchWindow window = searchWindowProvider.getSearchWindow(history, image.getWidth(), image.getHeight(), now);
            BufferedImage searchImage = image.getSubimage(window.getX(), window.getY(), window.getWidth(), window.getHeight());

            //Pre-Process the new image with algorithm
            for (PreProcessor preprocessor : preprocessors) {
                searchImage = preprocessor.process(searchImage);
            }

            //Search the target.
            //Build color histogram or regions.
            //Compare histograms or regions.
            final Map<Color, SortedSet<Segment>> currentSegementsMap = segmentator.segment(searchImage);

            //Modo de comparar primero buscar los segmentos del target en la ventana de busqueda
            //Una vez encontrado encontrado comparar las distancias entre los centrodies (COG)
            // Si las distancias no dan probar con otro candidato de ser posible.
            final Iterator<Segment> targetSegmentIterator = targetSegments.iterator();
            final Segment mainTarget = targetSegmentIterator.next();
            final SortedSet<Segment> currentCandidateSegment = currentSegementsMap.get(mainTarget.getColor());
            float currentZoom;

            int objectWeightX;
            int objectWeightY;
            int totalPointsCount;
            final TrackingObject previous = history.peek();
            final Map<Integer, Segment> previousSegments = previous.getAdditionalData(SEGMENTS_BY_ID);
            final Map<Integer, Segment> currentSegments = new HashMap<Integer, Segment>();
            if (currentCandidateSegment != null) {
                Segment mainCandidate = currentCandidateSegment.first();
                final Point mainCandidateCog = mainCandidate.getCog();
                drawPoint(searchImage, mainCandidateCog, mainTarget.getId());
                currentCandidateSegment.remove(mainCandidate);
                objectWeightX = mainCandidateCog.x * mainTarget.getSize();
                objectWeightY = mainCandidateCog.y * mainTarget.getSize();
                currentZoom = getZoom(mainTarget, mainCandidate);// * mainTarget.getSize();
                totalPointsCount = mainTarget.getSize();
                for (; targetSegmentIterator.hasNext();) {
                    Segment target = targetSegmentIterator.next();
                    final SortedSet<Segment> candidates = currentSegementsMap.get(target.getColor());
                    if (candidates != null && !candidates.isEmpty()) {
                        Segment candidate = searchTarget(target, previousSegments.get(target.getId()), candidates, mainCandidateCog);
                        if (candidate != null) {
                            currentSegments.put(target.getId(), candidate);
                            drawPoint(searchImage, candidate.getCog(), target.getId());
                            candidates.remove(candidate);
                            objectWeightX += candidate.getCog().x * target.getSize();
                            objectWeightY += candidate.getCog().y * target.getSize();
                            totalPointsCount += target.getSize();
//                            currentZoom += getZoom(target, candidate) * target.getSize();
                        }
                    }
                }

                final TrackingObject initial = initialElement;
                float zoomFactor = currentZoom;
                int count = 1;
                for (TrackHistory.TrackerHistoryElement<TrackingObject> element : history) {
                    final Float zoom = element.getModel().getAdditionalData(ZOOM);
                    if (zoom != null) {
                        zoomFactor += zoom;
                        count++;
                    }
                }
                zoomFactor = zoomFactor / (count);
                int width;
                int height;
                if (zoom_enable) {
                    width = (int) (initial.getWidth() * zoomFactor);
                    height = (int) (initial.getHeight() * zoomFactor);
                } else {
                    width = initial.getWidth();
                    height = initial.getHeight();
                    zoomFactor = 1;
                }
                final int newX = objectWeightX / totalPointsCount;
                final int newY = objectWeightY / totalPointsCount;

                final TrackingObjectImpl objectFound = new TrackingObjectImpl(window.getX() + newX - (width / 2), (window.getY()) + newY - (height / 2), width, height);
                objectFound.addAdditionalData(ZOOM, currentZoom);
                objectFound.addAdditionalData(SEGMENTS_BY_ID, currentSegments);
                EventListenerManger.getInstance().objectFound(objectFound,targetName, zoomFactor,image);
                history.push(objectFound, now);
            }

            image.getGraphics().drawImage(searchImage, window.getX(), window.getY(), null);
        }
        //Add to tracking history.

    }

    /*
   Search for the best candidate for the given target
    */
    private Segment searchTarget(Segment target, Segment previous, SortedSet<Segment> candidates, Point mainCandidateCog) {
        //Aca hay que proponer una formula que tenga en cuenta la posicion (distancia con el centro del main y posicion anterior) y el tamano
        final Double targetDistance = targetDistances.get(target.getId());
        Segment betterTarget = null;


        double difference = 0;
        for (Segment candidate : candidates) {
            double currentDiference;
            final double distance;
            final double size;
            final double shape;
            if (previous == null) {
                distance = getDistance(mainCandidateCog, candidate.getCog()) / targetDistance;
                size = candidate.getSize() / target.getSize();
                shape = Math.abs(candidate.getWidth() / candidate.getHeight() - target.getWidth() / target.getHeight());
                if (((size <= 1.5 && size >= 0.5) || shape <= 1.5)) {
                    currentDiference = distance + size * 10 + shape * 10;
                } else {
                    continue;
                }
            } else {
                distance = getDistance(previous.getCog(), candidate.getCog());
                size = candidate.getSize() / previous.getSize();
                shape = Math.abs(candidate.getWidth() / candidate.getHeight() - previous.getWidth() / previous.getHeight());
                if (((size <= 1.1 && size >= 0.9) || shape <= 1)) {
                    currentDiference = distance;
                } else {
                    continue;
                }
            }
            if (betterTarget == null || currentDiference < difference) {
                betterTarget = candidate;
                difference = currentDiference;


            }
        }


        return betterTarget;
    }

    private void drawPoint(BufferedImage searchImage, Point cog, int id) {
        Graphics graphics = searchImage.getGraphics();
        graphics.setColor(java.awt.Color.BLACK);
        graphics.drawLine(cog.x - 3, cog.y, cog.x + 3, cog.y);
        graphics.drawLine(cog.x, cog.y - 3, cog.x, cog.y + 3);
        graphics.drawString(id + "", cog.x, cog.y);

    }

    public BufferedImage setTarget(BufferedImage image, int x, int y, String name, boolean zoom) {
        this.targetName = name;
        history = new TrackHistory<TrackingObject>(HISTORY_SIZE);
        BufferedImage processedImage = image;
        //Pre-Process the new image with algorithm
        for (PreProcessor preprocessor : preprocessors) {
            processedImage = preprocessor.process(processedImage);
        }

        segmentator.segment(processedImage);

        Map<Integer, Segment> segmentsMap = segmentator.getSegmentsMap();
        this.targetDistances = new HashMap<Integer, Double>();
        this.targetSegments = segmentator.getSegments();

        Point mainCog = targetSegments.first().getCog();
        for (Segment targetSegment : targetSegments) {
            final Point cog = targetSegment.getCog();
            targetDistances.put(targetSegment.getId(), getDistance(mainCog, cog));
            drawPoint(processedImage, cog, targetSegment.getId());
        }

        initialElement = new TrackingObjectImpl(x, y, image.getWidth(), image.getHeight());
        initialElement.addAdditionalData(ZOOM, 1f);
        initialElement.addAdditionalData(SEGMENTS_BY_ID, segmentsMap);
        history.push(initialElement, System.currentTimeMillis());
        zoom_enable = zoom;
        //This must be the last
        this.targetImage = processedImage;

        EventListenerManger.getInstance().registerNewObjectToTrack(name);

        return processedImage;
    }

    private double getDistance(Point main, Point cog) {
        return MathFunctions.euclideanDistance(new double[]{cog.x, cog.y}, new double[]{main.x, main.y});
    }

    public void removeTarget() {
        targetImage = null;
    }

    public static float getZoom(Segment original, Segment newOne) {
        int originalWidth = original.getWidth();
        int originalHeight = original.getHeight();

        double originalRatio = (double) originalWidth / originalHeight;

        int widthVariation = originalWidth - newOne.getWidth();
        int heightVariation = originalHeight - newOne.getHeight();

        int modelWidth;

        if (Math.abs(widthVariation) / originalWidth <= Math.abs(heightVariation) / originalHeight) {
//        if (Math.abs(widthVariation)  <= Math.abs(heightVariation)) {
            modelWidth = newOne.getWidth();
        } else {
            modelWidth = (int) (newOne.getHeight() * originalRatio);
        }

        return ((float) modelWidth / originalWidth);
    }
}
